import torch

class ModelConfig(object):
    def __init__(self, feas_size=80, \
                 num_classes=40, \
                 hidden_size=512, \
                 num_layers=4, \
                 model_type="blstm", \
                 post_rnn_mode="concat"):

        self.num_classes = num_classes
        self.feas_size = feas_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.model_type = model_type
        self.post_rnn_mode = post_rnn_mode


class TrainConfig(object):
    def __init__(self, save_file="blstm_lr0.001.pt", \
                 scp_path="/home/changhengyi/zps_rnn/train.scp", \
                 batch_size=256, \
                 lr=0.001, \
                 num_epochs=50, \
                 njobs=10):

        self.save_file = save_file
        self.scp_path = scp_path
        self.batch_size = batch_size
        self.lr = lr
        self.num_epochs = num_epochs
        self.njobs = njobs

class TestConfig(TrainConfig):
    def __init__(self):
        super(TestConfig, self).__init__()
        self.scp_path = self.scp_path.replace("train", "test")
        # config gpu or not
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")


